package com.luobi.rabbitmq.interceptor;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.amqp.rabbit.config.SimpleRabbitListenerContainerFactory;
import org.springframework.amqp.rabbit.connection.CachingConnectionFactory;
import org.springframework.amqp.rabbit.connection.ConnectionFactory;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.boot.autoconfigure.amqp.SimpleRabbitListenerContainerFactoryConfigurer;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Map;
import java.util.UUID;

import static com.luobi.rabbitmq.constants.TraceIdConstants.CONTEXT_TRACE_ID;
import static com.luobi.rabbitmq.constants.TraceIdConstants.USER_REQ_ID_HEADER_KEY;

@Configuration
public class RabbitInterceptor {

    @Bean(name = "rabbitListenerContainerFactory")
    @ConditionalOnProperty(prefix = "spring.rabbitmq.listener", name = "type", havingValue = "simple", matchIfMissing = true)
    public SimpleRabbitListenerContainerFactory simpleRabbitListenerContainerFactory(SimpleRabbitListenerContainerFactoryConfigurer configurer, ConnectionFactory connectionFactory) {
        SimpleRabbitListenerContainerFactory factory = new SimpleRabbitListenerContainerFactory();
        // 消息接收之前加拦截处理，每次接收消息都会调用 setAfterReceivePostProcessors
        factory.setAfterReceivePostProcessors(message -> {
            Map<String, Object> header = message.getMessageProperties().getHeaders();
            if (header.containsKey(USER_REQ_ID_HEADER_KEY)) {  // 从消息头中获取traceId，如果存在，就设置到MDC中
                MDC.put(CONTEXT_TRACE_ID, header.get(USER_REQ_ID_HEADER_KEY).toString());
            }
            return message;
        });
        configurer.configure(factory, connectionFactory);
        return factory;
    }

    @Bean
    public RabbitTemplate rabbitTemplate(CachingConnectionFactory factory) {
        RabbitTemplate rabbitTemplate = new RabbitTemplate(factory);
        // 消息发送之前加一个拦截器，每次发送都会调用 setBeforePublishPostProcessors
        rabbitTemplate.setBeforePublishPostProcessors(message -> {
            String traceId = MDC.get(CONTEXT_TRACE_ID);  // 从MDC中获取traceId
            if (StringUtils.isBlank(traceId)) {
                traceId = StringUtils.replace(UUID.randomUUID().toString(), "-", "");
            }
            message.getMessageProperties().getHeaders().put(USER_REQ_ID_HEADER_KEY, traceId);  // 设置traceId到消息头中
            return message;
        });
        return rabbitTemplate;
    }

}
